import os
import numpy as np
import tensorflow as tf

class ResNet:
    def __init__(self, name, n, n_classes=10, batch_size=50, drop_set=[]):
        self.name = name
        self.n = n
        self.drop_set = drop_set
        self.n_classes = n_classes
        self.batch_size = batch_size

        self.bi = [0] * 3
        for i in range(self.n * 3):
            if i in drop_set:
                continue
            else:
                self.bi[i//self.n] += 1 #i//self.n的结果是0，1，2。网络中一共有3个block，计算出每个block一共有多少层
                                        #之所以要计算是因为剪枝之后模型就不规则了，必须计算一下每个block各有多少层
        
        # print(self.bi)
        
        self.ss_size = self.bi[0]*16 + self.bi[1]*32 + self.bi[2]*64# ss_size是每个模型中filter的总数
        self.data_shape = (32, 32, 3)

        self._define_inputs()
        self._build_graph()
        self._init_session()

        self.var_count()

    def _init_session(self):
        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True
        self.sess = tf.Session(config=config)
        self.sess.run(tf.global_variables_initializer())
        self.saver = tf.train.Saver(max_to_keep=10)
    
    def _define_inputs(self):
        shape = [self.batch_size]
        shape.extend(self.data_shape)
        self.images = tf.placeholder(
            tf.float32,
            shape=shape,
            name='input_images'
        )
        self.labels = tf.placeholder(
            tf.int32,
            shape=[self.batch_size, self.n_classes],
            name='labels'
        )
        self.learning_rate = tf.placeholder(
            tf.float32,
            shape=[],
            name='learning_rate'
        )
        self.momentum = tf.placeholder(
            tf.float32,
            shape=[],
            name='momentum'
        )
        self.keep_prob = tf.placeholder(
            tf.float32,
            shape=[],
            name='keep_prob'
        )
        self.weight_decay = tf.placeholder(
            tf.float32,
            shape=[],
            name='weight_decay'
        )
        self.block_ss = tf.placeholder(
            tf.float32,
            shape=[self.n*3],
            name='block_ss'
        )
        self.filter_ss = tf.placeholder(
            tf.float32,
            shape=[self.ss_size],
            name='filter_ss'
        )
        self.is_training = tf.placeholder(tf.bool, shape=[])
    
    def _build_graph(self):
        with tf.variable_scope(self.name):
            l = self.conv2d('conv0', self.images, 16)
            l = self.BNReLU(l)

            cnt = 0
            index = 0

            l = self.residual('res1.0', l, index, cnt, first=True)
            cnt += 16
            index += 1
            for k in range(1, self.n):
                if index in self.drop_set:
                    index += 1
                    continue
                l = self.residual('res1.{}'.format(k), l, index, cnt)
                cnt += 16
                index += 1
            
            # print(cnt, index)

            l = self.residual('res2.0', l, index, cnt, increase_dim=True)
            cnt += 32
            index += 1
            for k in range(1, self.n):
                if index in self.drop_set:
                    index += 1
                    continue
                l = self.residual('res2.{}'.format(k), l, index, cnt)
                cnt += 32
                index += 1
            
            # print(cnt, index)

            l = self.residual('res3.0', l, index, cnt, increase_dim=True)
            cnt += 64
            index += 1
            for k in range(1, self.n):
                if index in self.drop_set:
                    index += 1
                    continue
                l = self.residual('res3.{}'.format(k), l, index, cnt)
                cnt += 64
                index += 1
            
            # print(cnt, index)
            
            l = self.BNReLU(l)
            l = self.globalavgpooling('gap', l)
        
        logits = self.fullyconnected('linear', l, self.n_classes)
        prediction = tf.nn.softmax(logits)
        cross_entropy = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(
            logits=logits, labels=self.labels
        ))
        self.cross_entropy = cross_entropy
        l2_loss = tf.add_n(
            [tf.nn.l2_loss(var) for (var) in tf.trainable_variables()]
        )

        optimizer = tf.train.MomentumOptimizer(self.learning_rate, self.momentum)

        self.train_step = optimizer.minimize(cross_entropy + l2_loss*self.weight_decay)

        correct_prediction = tf.equal(
            tf.argmax(prediction, 1),
            tf.argmax(self.labels, 1)
        )
        self.accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
    
    def residual(self, name, l, index, cnt, increase_dim=False, first=False):
        shape = l.get_shape().as_list()
        in_channel = shape[-1]

        factor = self.block_ss[index]

        if increase_dim:
            out_channel = in_channel * 2
            stride1 = [1, 2, 2, 1]
        else:
            out_channel = in_channel
            stride1 = [1, 1, 1, 1]
        
        with tf.variable_scope(self.name+'/'+name):
            b1 = self.BNReLU(l)
            c1 = self.conv2d('conv1', b1, out_channel, kernel_size=3, strides=stride1)
            c1 = self.BNReLU(c1)
            c1 *= self.filter_ss[cnt:cnt+out_channel]
            # print(cnt, cnt+out_channel)
            c2 = self.conv2d('conv2', c1, out_channel, kernel_size=3)
            if increase_dim:
                l = self.avgpooling('pool', l, 2)
                l = tf.pad(l, [[0, 0], [0, 0], [0, 0], [in_channel//2, in_channel//2]])
            l = c2 * factor + l
        return l
    
    def fullyconnected(self, name, _input, out_channel):
        shape = _input.get_shape()
        if len(shape) == 4:
            size = shape[1].value * shape[2].value * shape[3].value
        else:
            size = shape[-1].value
        
        with tf.variable_scope(name):
            w = tf.get_variable(
                name='weights',
                shape=[size, out_channel],
                initializer=tf.contrib.layers.variance_scaling_initializer()
            )

            b = tf.get_variable('biases', initializer=tf.constant(0.0, shape=[out_channel]))

            flat_x = tf.reshape(_input, [-1, size])

            output = tf.nn.bias_add(tf.matmul(flat_x, w), b)
            return output
    
    def BNReLU(self, l):
        l = tf.contrib.layers.batch_norm(
            l,
            scale=True,
            is_training=self.is_training,
            updates_collections=None
        )
        l = tf.nn.relu(l)
        return l
    
    def conv2d(self, name, _input, out_channel, kernel_size=3, strides=[1, 1, 1, 1], padding='SAME'):
        in_channel = int(_input.get_shape()[-1])
        with tf.variable_scope(name):
            kernel = tf.get_variable(
                name='kernel',
                shape=[kernel_size, kernel_size, in_channel, out_channel],
                initializer=tf.variance_scaling_initializer(scale=2.0, mode='fan_out')
            )
            output = tf.nn.conv2d(_input, kernel, strides, padding)
            return output
    
    def avgpooling(self, name, _input, kernel):
        kernel_size = [1, kernel, kernel, 1]
        strides = [1, kernel, kernel, 1]
        padding = 'VALID'
        with tf.variable_scope(name):
            output = tf.nn.avg_pool(_input, kernel_size, strides, padding)
            return output
    
    def globalavgpooling(self, name, _input):
        last_pool_kernel = int(_input.get_shape()[-2])
        output = self.avgpooling(name, _input, last_pool_kernel)
        total_features = int(output.get_shape()[-1])
        output = tf.reshape(output, [-1, total_features])
        return output

    def model_test(self, data_provider, hyper_param):
        num_test_examples = data_provider.test.num_examples
        total_acc = []
        total_los = []

        for _ in range(num_test_examples // self.batch_size):
            batch = data_provider.test.next_batch(self.batch_size)
            feed_dict = {
                self.images: batch[0],
                self.labels: batch[1],
                self.keep_prob: 1.0,
                self.weight_decay: hyper_param['weight_decay'],
                self.block_ss: hyper_param['block_ss'],
                self.filter_ss: hyper_param['filter_ss'],
                self.is_training: False
            }
            fetches = [self.cross_entropy, self.accuracy]
            loss, accuracy = self.sess.run(fetches, feed_dict=feed_dict)
            total_los.append(loss)
            total_acc.append(accuracy)
        
        mean_acc = np.mean(total_acc)
        mean_los = np.mean(total_los)
        
        return mean_acc, mean_los
    
    def val_test(self, data_provider, hyper_param):
        num_test_examples = data_provider.validation.num_examples
        total_acc = []
        total_los = []

        for _ in range(num_test_examples // self.batch_size):
            batch = data_provider.validation.next_batch(self.batch_size)
            feed_dict = {
                self.images: batch[0],
                self.labels: batch[1],
                self.keep_prob: 1.0,
                self.weight_decay: hyper_param['weight_decay'],
                self.block_ss: hyper_param['block_ss'],
                self.filter_ss: hyper_param['filter_ss'],
                self.is_training: False
            }
            fetches = [self.cross_entropy, self.accuracy]
            loss, accuracy = self.sess.run(fetches, feed_dict=feed_dict)
            total_los.append(loss)
            total_acc.append(accuracy)
        
        mean_acc = np.mean(total_acc)
        mean_los = np.mean(total_los)
        
        return mean_acc, mean_los

    def train_single_epoch(self, data_provider, hyper_param):
        num_examples = data_provider.train.num_examples
        total_acc = []
        total_los = []
        for i in range(1, (num_examples//self.batch_size) + 1):
            batch = data_provider.train.next_batch(self.batch_size)
            images, labels = batch
            feed_dict = {
                self.images: images,
                self.labels: labels,
                self.momentum: hyper_param['momentum'],
                self.keep_prob: hyper_param['keep_prob'],
                self.learning_rate: hyper_param['learning_rate'],
                self.weight_decay: hyper_param['weight_decay'],
                self.block_ss: hyper_param['block_ss'],
                self.filter_ss: hyper_param['filter_ss'],
                self.is_training: True
            }
            fetches = [self.train_step, self.cross_entropy, self.accuracy]
            result = self.sess.run(fetches, feed_dict=feed_dict)
            _, loss, accuracy = result
            total_los.append(loss)
            total_acc.append(accuracy)

        mean_acc = np.mean(total_acc)
        mean_los = np.mean(total_los)

        return mean_acc, mean_los
    
    def var_count(self):
        total_parameters = 0
        for variable in tf.trainable_variables():
            shape = variable.get_shape()
            variable_parameters = 1
            for dim in shape:
                variable_parameters *= dim.value
            total_parameters += variable_parameters
        print('Total-parameters:%d' % (total_parameters))